#include <iostream>
#include <array>
#include <ctime>

#define INPUT_SIZE 2000


#ifdef WAFER_OPENSSL
#include <openssl/sha.h>
#endif

#include "container.h"

#ifndef WAFER_OPENSSL
extern "C"  {
    void CUDA(MemRef<unsigned char, 1>* input, int n, MemRef<unsigned char, 1>* output, int N);
}

void SHA1_MULTI_OP(MemRef<unsigned char, 1>* input, int n, MemRef<unsigned char, 1>* output, int N) {
    CUDA(input, n, output, N);
}
#endif

int n = 20;
int N = 100;
uint8_t exampleInput[INPUT_SIZE];
uint8_t exampleOutput[200000];
MemRef<unsigned char, 1> input{exampleInput, INPUT_SIZE};
MemRef<unsigned char, 1> output{exampleOutput, 200000};


void printResult(uint8_t* hash, std::string head, bool printHead) {
    // Print the output.
    if (printHead) {
      std::cout << "--------------------------------------------------------"
                  << std::endl;
      std::cout << head << std::endl;
    }
    std::cout << "[ ";
    for (size_t i = 0; i < 20; i++) {
     printf("%02x ", hash[i]);
    }
    std::cout << "]" << std::endl;
}

void prepareInput() {
  for(int i = 0; i < INPUT_SIZE; i++) {
    exampleInput[i] = uint8_t(i % 26) + 'a';
  }
}

void generateResultGPUSha1() {
#ifdef WAFER_OPENSSL
  std::cout << "--------------------------------------------------------\nOpenSSL MD5\n";

  const clock_t begin_time = clock();
  for (int i = 0; i < N; i++) {
    SHA_CTX context;
    SHA1_Init(&context);
    SHA1_Update(&context, exampleInput + (i * n), n);
    SHA1_Final(exampleOutput + (i * SHA_DIGEST_LENGTH), &context);

    printResult(exampleOutput + (i * SHA_DIGEST_LENGTH), "MLIR SHA1", i == 0);
  }
  float duration = float(clock() - begin_time) / CLOCKS_PER_SEC;
  std::cout << "Calculate duration: " << std::fixed << duration << std::endl;
  std::cout << "--------------------------------------------------------"
            << std::endl;

#else
  const clock_t begin_time = clock();

  #pragma wafer parallel
  SHA1_MULTI_OP(&input, n, &output, N);

  float duration = float(clock() - begin_time) / CLOCKS_PER_SEC;

  std::cout << "--------------------------------------------------------\nMLIR SHA1\n";
  for (int i = 0; i < N; i++) {
    printResult(exampleOutput + i * 20, "MLIR SHA1", i == 0);
  }
  std::cout << "Calculate duration: " << std::fixed << duration << std::endl;
  std::cout << "--------------------------------------------------------"
            << std::endl;
#endif
}

int main() {
  prepareInput();
  generateResultGPUSha1();
  return 0;
}

